
use super::solution::Solution;


use std::collections::HashMap;
struct UnionSet {
    parent: Vec<i32>,
    rank :Vec<i32>
}

impl UnionSet {
    pub fn new(len:usize) -> UnionSet {
        let mut parent = Vec::with_capacity(len);
        for i in 0..len{
            parent.push(i as i32);
        }
        UnionSet {
            parent,
            rank: [1].repeat(len),
        }
    }

    pub fn find(&mut self,i :i32) -> i32{
        if self.parent[i as usize] == i{
            return i ;
        }
        let p=self.find(self.parent[i as usize]);
        self.parent[i as usize] = p;
        p
    }
    pub fn union(&mut self,i:i32,j :i32){
        let (pi,pj) = (self.find(i), self.find(j));
        let (ri,rj)=(self.rank[pi as usize], self.rank[pj as usize]);
        if ri==rj{
            self.parent[pj as usize] = pi;
            self.rank[pi as usize] += 1;
        }else if ri>=rj{
            self.parent[pj as usize] = pi;
        }else{
            self.parent[pi as usize] = pj;
        }
    }
}
impl Solution {
    pub fn count_pairs(n: i32,mut edges: Vec<Vec<i32>>) -> i64 {
        let mut uf = UnionSet::new(n as usize);
        let mut ans =0i64;
        for p in edges{
            uf.union(p[0],p[1]);
        }
        let mut map=HashMap::with_capacity(n as usize);
        for i in 0..n as usize{
            map.entry(uf.find(i as i32)).
                and_modify(|e| *e += 1).
                or_insert(1);
        }
        // println!("{:?}",map);
        for (_,v) in map.iter(){
            let v:i64=*v;
            ans +=v*(n as i64-v);
        }
        
        ans/2
    }
}